import sys
import os
import cv2
import numpy as np
from PIL import Image, ImageQt
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                            QHBoxLayout, QPushButton, QLabel, QFileDialog, 
                            QComboBox, QMessageBox, QProgressBar, QGroupBox,
                            QSlider, QGridLayout, QSplitter, QCheckBox)
from PyQt5.QtGui import QPixmap, QImage, QFont
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from ultralytics import YOLO
import matplotlib
matplotlib.use('Agg')  # 使用Agg后端，避免显示问题
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import time
from matplotlib import font_manager

# 尝试设置中文字体
def set_chinese_font():
    """尝试设置可用的中文字体"""
    # 获取系统中所有可用字体
    available_fonts = {f.name for f in font_manager.fontManager.ttflist}
    print(f"可用字体: {available_fonts}")  # 打印可用字体，用于调试
    
    # 常用中文字体列表
    chinese_fonts = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC", "Microsoft YaHei", "SimSun"]
    
    for font in chinese_fonts:
        if font in available_fonts:
            plt.rcParams["font.family"] = font
            print(f"已设置中文字体: {font}")
            return True
    
    # 如果没有找到任何中文字体，则使用默认字体
    print("警告: 未找到可用的中文字体，图表将使用英文标签")
    
    # 设置英文标签
    plt.rcParams["font.family"] = ["DejaVu Sans", "sans-serif"]
    
    # 替换所有中文标签为英文
    global RIPENESS_MAPPING
    RIPENESS_MAPPING = {
        "成熟": "Ripe",
        "半成熟": "Half-ripe",
        "未成熟": "Unripe",
        "过熟": "Overripe",
        "苹果": "Apple",
        "香蕉": "Banana",
        "橙子": "Orange",
        "草莓": "Strawberry",
        "葡萄": "Grape"
    }
    
    return False

# 初始化成熟度映射
RIPENESS_MAPPING = {
    "成熟": "成熟",
    "半成熟": "半成熟",
    "未成熟": "未成熟",
    "过熟": "过熟",
    "苹果": "苹果",
    "香蕉": "香蕉",
    "橙子": "橙子",
    "草莓": "草莓",
    "葡萄": "葡萄"
}

# 设置中文字体
set_chinese_font()

class DetectionThread(QThread):
    """
    检测线程，用于在后台运行YOLOv8检测，防止界面卡顿
    """
    update_signal = pyqtSignal(np.ndarray, list)
    finished_signal = pyqtSignal()
    progress_signal = pyqtSignal(int)
    
    def __init__(self, model, image_path, fruit_type, confidence_threshold, image_preprocessing=False):
        super().__init__()
        self.model = model
        self.image_path = image_path
        self.fruit_type = fruit_type
        self.confidence_threshold = confidence_threshold
        self.image_preprocessing = image_preprocessing
        
        # 改进的COCO类别映射
        self.coco_classes = [
            "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", 
            "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", 
            "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", 
            "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", 
            "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", 
            "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", 
            "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", 
            "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", 
            "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", 
            "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
        ]
        
        # 改进的水果类别映射
        self.fruit_mapping = {
            "苹果": "apple",
            "香蕉": "banana",
            "橙子": "orange",
            "草莓": "strawberry",  # 尝试添加草莓
            "葡萄": "grape"        # 尝试添加葡萄
        }
        
    def run(self):
        try:
            print(f"开始检测: {self.image_path}, 水果类型: {self.fruit_type}, 置信度阈值: {self.confidence_threshold}")
            
            # 模拟进度
            self.progress_signal.emit(10)
            time.sleep(0.2)
            
            # 加载图像
            image = cv2.imread(self.image_path)
            if image is None:
                print(f"错误: 无法加载图像 {self.image_path}")
                return
            
            # 图像预处理
            if self.image_preprocessing:
                print("应用图像预处理...")
                image = self.preprocess_image(image)
            
            self.progress_signal.emit(30)
            time.sleep(0.2)
            
            # 运行检测
            results = self.model.predict(
                source=image,
                conf=self.confidence_threshold,
                classes=self.get_class_indices(),
                verbose=False
            )
            
            print(f"检测完成，发现 {len(results[0].boxes) if results[0].boxes is not None else 0} 个目标")
            
            self.progress_signal.emit(60)
            time.sleep(0.2)
            
            # 处理结果
            processed_image = results[0].plot() if results[0].boxes is not None and len(results[0].boxes) > 0 else image.copy()
            detections = self.process_detections(results[0])
            
            print(f"处理完成，识别出 {len(detections)} 个水果")
            
            self.progress_signal.emit(90)
            time.sleep(0.2)
            
            # 发送更新信号
            self.update_signal.emit(processed_image, detections)
            
        except Exception as e:
            print(f"检测过程中出错: {str(e)}")
        finally:
            self.progress_signal.emit(100)
            self.finished_signal.emit()
    
    def preprocess_image(self, image):
        """图像预处理以提高检测效果"""
        # 转换为RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        
        # 高斯模糊减少噪声
        image = cv2.GaussianBlur(image, (5, 5), 0)
        
        # 直方图均衡化增强对比度
        ycrcb = cv2.cvtColor(image, cv2.COLOR_RGB2YCrCb)
        channels = cv2.split(ycrcb)
        channels[0] = cv2.equalizeHist(channels[0])
        ycrcb = cv2.merge(channels)
        image = cv2.cvtColor(ycrcb, cv2.COLOR_YCrCb2RGB)
        
        # 转回BGR
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        
        return image
    
    def get_class_indices(self):
        """获取所选水果类型的类别索引"""
        if self.fruit_type == "所有水果":
            return None  # 检测所有类别
        
        # 查找水果在COCO类别中的索引
        english_name = self.fruit_mapping.get(self.fruit_type)
        if english_name and english_name in self.coco_classes:
            index = self.coco_classes.index(english_name)
            print(f"筛选类别: {self.fruit_type} ({english_name}), 索引: {index}")
            return [index]
        else:
            print(f"警告: 水果类型 '{self.fruit_type}' 在COCO数据集中没有找到精确匹配，将检测所有类别")
            return None  # 返回None以检测所有类别
    
    def process_detections(self, results):
        """处理检测结果，提取水果和成熟度信息"""
        detections = []
        
        if results.boxes is not None and len(results.boxes) > 0:
            for i, box in enumerate(results.boxes):
                # 获取边界框坐标
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                
                # 获取类别名称和置信度
                class_id = int(box.cls)
                confidence = float(box.conf)
                
                # 检查类别ID是否在COCO类别列表范围内
                if class_id < len(self.coco_classes):
                    coco_name = self.coco_classes[class_id]
                    
                    # 映射COCO类别到我们的水果类型
                    fruit_mapping = {
                        "apple": "苹果",
                        "banana": "香蕉",
                        "orange": "橙子",
                        "strawberry": "草莓",  # 添加草莓支持
                        "grape": "葡萄"        # 添加葡萄支持
                    }
                    
                    fruit_name = fruit_mapping.get(coco_name, "未知水果")
                    
                    # 如果选择了特定水果类型，且当前检测结果不匹配，则跳过
                    if self.fruit_type != "所有水果" and fruit_name != self.fruit_type:
                        continue
                    
                    if fruit_name != "未知水果":
                        print(f"检测到 {fruit_name}, 置信度: {confidence:.2f}, 边界框: ({x1}, {y1}, {x2}, {y2})")
                        
                        # 提取水果区域
                        fruit_region = results.orig_img[y1:y2, x1:x2]
                        
                        # 分析成熟度
                        ripeness, color_info = self.analyze_ripeness(fruit_name, fruit_region)
                        
                        # 添加到检测结果列表
                        detections.append({
                            "box": (x1, y1, x2, y2),
                            "fruit_name": fruit_name,
                            "confidence": confidence,
                            "ripeness": ripeness,
                            "color_info": color_info
                        })
                    else:
                        print(f"忽略检测结果: {coco_name}, 置信度: {confidence:.2f}")
                else:
                    print(f"警告: 未知类别ID {class_id}, 置信度: {confidence:.2f}")
        
        return detections
    
    def analyze_ripeness(self, fruit_name, fruit_region):
        """分析水果成熟度"""
        if fruit_region.size == 0:
            return "无法分析", {}
        
        # 转换为HSV颜色空间
        hsv_image = cv2.cvtColor(fruit_region, cv2.COLOR_BGR2HSV)
        
        # 根据不同水果类型使用不同的成熟度判断方法
        if fruit_name == "苹果":
            return self.analyze_apple_ripeness(hsv_image)
        elif fruit_name == "香蕉":
            return self.analyze_banana_ripeness(hsv_image)
        elif fruit_name == "橙子":
            return self.analyze_orange_ripeness(hsv_image)
        elif fruit_name == "草莓":
            return self.analyze_strawberry_ripeness(hsv_image)
        elif fruit_name == "葡萄":
            return self.analyze_grape_ripeness(hsv_image)
        else:
            return "未知水果类型", {}
    
    def analyze_apple_ripeness(self, hsv_image):
        """分析苹果成熟度"""
        # 红色区域 (两个范围，因为红色在HSV中是分段的)
        lower_red1 = np.array([0, 100, 100])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([160, 100, 100])
        upper_red2 = np.array([180, 255, 255])
        
        # 绿色区域
        lower_green = np.array([40, 100, 100])
        upper_green = np.array([80, 255, 255])
        
        # 创建掩码
        mask_red1 = cv2.inRange(hsv_image, lower_red1, upper_red1)
        mask_red2 = cv2.inRange(hsv_image, lower_red2, upper_red2)
        mask_red = cv2.bitwise_or(mask_red1, mask_red2)
        mask_green = cv2.inRange(hsv_image, lower_green, upper_green)
        
        # 计算红色和绿色像素比例
        total_pixels = hsv_image.shape[0] * hsv_image.shape[1]
        red_pixels = cv2.countNonZero(mask_red)
        green_pixels = cv2.countNonZero(mask_green)
        
        red_percentage = red_pixels / total_pixels * 100
        green_percentage = green_pixels / total_pixels * 100
        
        # 判断成熟度
        if red_percentage > 60:
            return "成熟", {"红色比例": f"{red_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        elif green_percentage > 40:
            return "未成熟", {"红色比例": f"{red_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        else:
            return "半成熟", {"红色比例": f"{red_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
    
    def analyze_banana_ripeness(self, hsv_image):
        """分析香蕉成熟度"""
        # 黄色区域
        lower_yellow = np.array([20, 100, 100])
        upper_yellow = np.array([40, 255, 255])
        
        # 绿色区域
        lower_green = np.array([40, 100, 100])
        upper_green = np.array([80, 255, 255])
        
        # 棕色区域
        lower_brown = np.array([10, 100, 100])
        upper_brown = np.array([20, 255, 255])
        
        # 创建掩码
        mask_yellow = cv2.inRange(hsv_image, lower_yellow, upper_yellow)
        mask_green = cv2.inRange(hsv_image, lower_green, upper_green)
        mask_brown = cv2.inRange(hsv_image, lower_brown, upper_brown)
        
        # 计算像素比例
        total_pixels = hsv_image.shape[0] * hsv_image.shape[1]
        yellow_pixels = cv2.countNonZero(mask_yellow)
        green_pixels = cv2.countNonZero(mask_green)
        brown_pixels = cv2.countNonZero(mask_brown)
        
        yellow_percentage = yellow_pixels / total_pixels * 100
        green_percentage = green_pixels / total_pixels * 100
        brown_percentage = brown_pixels / total_pixels * 100
        
        # 判断成熟度
        if green_percentage > 40:
            return "未成熟", {"黄色比例": f"{yellow_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%", "棕色比例": f"{brown_percentage:.1f}%"}
        elif yellow_percentage > 60:
            return "成熟", {"黄色比例": f"{yellow_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%", "棕色比例": f"{brown_percentage:.1f}%"}
        elif brown_percentage > 20:
            return "过熟", {"黄色比例": f"{yellow_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%", "棕色比例": f"{brown_percentage:.1f}%"}
        else:
            return "半成熟", {"黄色比例": f"{yellow_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%", "棕色比例": f"{brown_percentage:.1f}%"}
    
    def analyze_orange_ripeness(self, hsv_image):
        """分析橙子成熟度"""
        # 橙色区域
        lower_orange = np.array([10, 100, 100])
        upper_orange = np.array([30, 255, 255])
        
        # 绿色区域
        lower_green = np.array([40, 100, 100])
        upper_green = np.array([80, 255, 255])
        
        # 创建掩码
        mask_orange = cv2.inRange(hsv_image, lower_orange, upper_orange)
        mask_green = cv2.inRange(hsv_image, lower_green, upper_green)
        
        # 计算像素比例
        total_pixels = hsv_image.shape[0] * hsv_image.shape[1]
        orange_pixels = cv2.countNonZero(mask_orange)
        green_pixels = cv2.countNonZero(mask_green)
        
        orange_percentage = orange_pixels / total_pixels * 100
        green_percentage = green_pixels / total_pixels * 100
        
        # 判断成熟度
        if orange_percentage > 70:
            return "成熟", {"橙色比例": f"{orange_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        elif green_percentage > 50:
            return "未成熟", {"橙色比例": f"{orange_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        else:
            return "半成熟", {"橙色比例": f"{orange_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
    
    def analyze_strawberry_ripeness(self, hsv_image):
        """分析草莓成熟度"""
        # 红色区域
        lower_red1 = np.array([0, 100, 100])
        upper_red1 = np.array([10, 255, 255])
        lower_red2 = np.array([160, 100, 100])
        upper_red2 = np.array([180, 255, 255])
        
        # 绿色区域 (草莓叶子)
        lower_green = np.array([40, 100, 100])
        upper_green = np.array([80, 255, 255])
        
        # 创建掩码
        mask_red1 = cv2.inRange(hsv_image, lower_red1, upper_red1)
        mask_red2 = cv2.inRange(hsv_image, lower_red2, upper_red2)
        mask_red = cv2.bitwise_or(mask_red1, mask_red2)
        mask_green = cv2.inRange(hsv_image, lower_green, upper_green)
        
        # 计算像素比例
        total_pixels = hsv_image.shape[0] * hsv_image.shape[1]
        red_pixels = cv2.countNonZero(mask_red)
        green_pixels = cv2.countNonZero(mask_green)
        
        red_percentage = red_pixels / total_pixels * 100
        green_percentage = green_pixels / total_pixels * 100
        
        # 判断成熟度
        if red_percentage > 80:
            return "成熟", {"红色比例": f"{red_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        elif red_percentage < 50:
            return "未成熟", {"红色比例": f"{red_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        else:
            return "半成熟", {"红色比例": f"{red_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
    
    def analyze_grape_ripeness(self, hsv_image):
        """分析葡萄成熟度"""
        # 紫色区域
        lower_purple = np.array([120, 100, 100])
        upper_purple = np.array([150, 255, 255])
        
        # 绿色区域
        lower_green = np.array([40, 100, 100])
        upper_green = np.array([80, 255, 255])
        
        # 创建掩码
        mask_purple = cv2.inRange(hsv_image, lower_purple, upper_purple)
        mask_green = cv2.inRange(hsv_image, lower_green, upper_green)
        
        # 计算像素比例
        total_pixels = hsv_image.shape[0] * hsv_image.shape[1]
        purple_pixels = cv2.countNonZero(mask_purple)
        green_pixels = cv2.countNonZero(mask_green)
        
        purple_percentage = purple_pixels / total_pixels * 100
        green_percentage = green_pixels / total_pixels * 100
        
        # 判断成熟度
        if purple_percentage > 70:
            return "成熟", {"紫色比例": f"{purple_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        elif green_percentage > 60:
            return "未成熟", {"紫色比例": f"{purple_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}
        else:
            return "半成熟", {"紫色比例": f"{purple_percentage:.1f}%", "绿色比例": f"{green_percentage:.1f}%"}


class MaturityVisualizer(FigureCanvas):
    """成熟度可视化图表"""
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        super().__init__(self.fig)
        self.setParent(parent)
        self.axes = self.fig.add_subplot(111)
        self.clear_plot()
        
    def clear_plot(self):
        """清空图表"""
        self.axes.clear()
        # 使用全局映射替换中文标签
        title = RIPENESS_MAPPING.get("水果成熟度分布", "Fruit Ripeness Distribution")
        x_label = RIPENESS_MAPPING.get("水果类型", "Fruit Type")
        y_label = RIPENESS_MAPPING.get("数量", "Count")
        
        self.axes.set_title(title)
        self.axes.set_xlabel(x_label)
        self.axes.set_ylabel(y_label)
        self.axes.set_ylim(0, 10)
        self.draw()
    
    def update_plot(self, detections):
        """更新图表"""
        self.clear_plot()
        
        if not detections:
            return
            
        # 统计每种水果的成熟度分布
        fruit_ripeness = {}
        
        for detection in detections:
            # 使用全局映射替换中文标签
            fruit_name = RIPENESS_MAPPING.get(detection["fruit_name"], detection["fruit_name"])
            ripeness = RIPENESS_MAPPING.get(detection["ripeness"], detection["ripeness"])
            
            if fruit_name not in fruit_ripeness:
                fruit_ripeness[fruit_name] = {"成熟": 0, "半成熟": 0, "未成熟": 0, "过熟": 0}
            
            fruit_ripeness[fruit_name][ripeness] += 1
        
        # 绘制图表
        fruits = list(fruit_ripeness.keys())
        if not fruits:
            return
            
        max_count = max(max(counts.values()) for counts in fruit_ripeness.values()) + 1
        self.axes.set_ylim(0, max_count)
        
        bar_width = 0.2
        positions = range(len(fruits))
        
        # 使用全局映射替换中文标签
        ripeness_labels = ["成熟", "半成熟", "未成熟", "过熟"]
        ripeness_labels = [RIPENESS_MAPPING.get(label, label) for label in ripeness_labels]
        
        for i, ripeness in enumerate(ripeness_labels):
            counts = [fruit_ripeness[fruit].get(ripeness, 0) for fruit in fruits]
            offset = i * bar_width - (3 * bar_width / 2)
            self.axes.bar([p + offset for p in positions], counts, width=bar_width, label=ripeness)
        
        self.axes.set_xticks(positions)
        self.axes.set_xticklabels(fruits, rotation=45)
        self.axes.legend()
        self.fig.tight_layout()
        self.draw()


class FruitRipenessDetector(QMainWindow):
    """水果成熟度检测主窗口"""
    def __init__(self):
        super().__init__()
        
        # 设置窗口标题和大小
        self.setWindowTitle("水果成熟度检测系统")
        self.resize(1200, 800)
        
        # 初始化模型和变量
        self.model = None
        self.current_image_path = None
        self.detections = []
        self.image_preprocessing = False
        
        # 创建中心部件和布局
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.main_layout = QVBoxLayout(self.central_widget)
        
        # 创建顶部控制区域
        self.create_control_area()
        
        # 创建分割器，将窗口分为左右两部分
        self.splitter = QSplitter(Qt.Horizontal)
        
        # 创建左侧图像显示区域
        self.create_image_display_area()
        
        # 创建右侧结果和统计区域
        self.create_results_area()
        
        # 将左右区域添加到分割器
        self.splitter.addWidget(self.image_group)
        self.splitter.addWidget(self.results_group)
        
        # 设置分割器初始大小
        self.splitter.setSizes([700, 500])
        
        # 将分割器添加到主布局
        self.main_layout.addWidget(self.splitter)
        
        # 加载模型
        self.load_model()
        
    def create_control_area(self):
        """创建顶部控制区域"""
        control_layout = QHBoxLayout()
        
        # 添加文件选择按钮
        self.select_image_btn = QPushButton("选择图片")
        self.select_image_btn.clicked.connect(self.select_image)
        control_layout.addWidget(self.select_image_btn)
        
        # 添加检测按钮
        self.detect_btn = QPushButton("开始检测")
        self.detect_btn.clicked.connect(self.detect_fruit)
        self.detect_btn.setEnabled(False)
        control_layout.addWidget(self.detect_btn)
        
        # 添加水果类型选择下拉框
        self.fruit_type_combo = QComboBox()
        self.fruit_type_combo.addItems(["所有水果", "苹果", "香蕉", "橙子", "草莓", "葡萄"])
        control_layout.addWidget(QLabel("水果类型:"))
        control_layout.addWidget(self.fruit_type_combo)
        
        # 添加置信度阈值滑块
        control_layout.addWidget(QLabel("置信度阈值:"))
        self.confidence_slider = QSlider(Qt.Horizontal)
        self.confidence_slider.setMinimum(1)
        self.confidence_slider.setMaximum(90)
        self.confidence_slider.setValue(30)  # 降低默认阈值以提高检测灵敏度
        self.confidence_slider.setTickInterval(10)
        self.confidence_slider.setTickPosition(QSlider.TicksBelow)
        self.confidence_slider.valueChanged.connect(self.update_confidence_label)
        control_layout.addWidget(self.confidence_slider)
        
        self.confidence_label = QLabel(f"{self.confidence_slider.value()/100:.2f}")
        control_layout.addWidget(self.confidence_label)
        
        # 添加图像预处理复选框
        self.preprocessing_checkbox = QCheckBox("图像预处理")
        self.preprocessing_checkbox.setChecked(self.image_preprocessing)
        self.preprocessing_checkbox.stateChanged.connect(self.toggle_preprocessing)
        control_layout.addWidget(self.preprocessing_checkbox)
        
        # 添加模型选择下拉框
        control_layout.addWidget(QLabel("模型:"))
        self.model_combo = QComboBox()
        self.model_combo.addItems(["yolov8n.pt", "yolov8s.pt", "yolov8m.pt"])
        self.model_combo.currentTextChanged.connect(self.change_model)
        control_layout.addWidget(self.model_combo)
        
        # 添加进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setTextVisible(True)
        self.progress_bar.setValue(0)
        control_layout.addWidget(self.progress_bar)
        
        control_layout.addStretch()
        
        self.main_layout.addLayout(control_layout)
    
    def update_confidence_label(self):
        """更新置信度标签"""
        self.confidence_label.setText(f"{self.confidence_slider.value()/100:.2f}")
    
    def toggle_preprocessing(self, state):
        """切换图像预处理选项"""
        self.image_preprocessing = state == Qt.Checked
        print(f"图像预处理已{'启用' if self.image_preprocessing else '禁用'}")
    
    def change_model(self, model_name):
        """更改使用的模型"""
        if self.model and model_name:
            reply = QMessageBox.question(
                self, "确认", f"是否要切换到{model_name}模型？这可能需要一些时间下载。",
                QMessageBox.Yes | QMessageBox.No, QMessageBox.No
            )
            
            if reply == QMessageBox.Yes:
                self.model = None
                self.load_model(model_name)
    
    def create_image_display_area(self):
        """创建图像显示区域"""
        self.image_group = QGroupBox("图像")
        image_layout = QVBoxLayout(self.image_group)
        
        self.image_label = QLabel("请选择一张图片")
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setMinimumSize(400, 400)
        self.image_label.setStyleSheet("border: 1px solid #ccc;")
        
        image_layout.addWidget(self.image_label)
    
    def create_results_area(self):
        """创建结果和统计区域"""
        self.results_group = QGroupBox("检测结果")
        results_layout = QVBoxLayout(self.results_group)
        
        # 创建结果表格区域
        self.results_table = QLabel("检测结果将显示在这里")
        self.results_table.setAlignment(Qt.AlignTop | Qt.AlignLeft)
        self.results_table.setWordWrap(True)
        
        # 创建图表区域
        self.visualizer = MaturityVisualizer(self, width=5, height=4)
        
        results_layout.addWidget(self.results_table)
        results_layout.addWidget(self.visualizer)
    
    def load_model(self, model_name="yolov8n.pt"):
        """加载YOLOv8模型"""
        try:
            # 显示加载消息
            self.results_table.setText(f"正在加载模型: {model_name}...")
            
            # 尝试加载模型
            self.model = YOLO(model_name)
            print(f"模型 {model_name} 加载成功")
            self.results_table.setText(f"模型 {model_name} 加载成功")
            
            # 更新模型下拉框
            current_index = self.model_combo.findText(model_name)
            if current_index >= 0:
                self.model_combo.setCurrentIndex(current_index)
            
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            QMessageBox.critical(self, "模型加载失败", f"加载模型时出错: {str(e)}")
            self.model = None
    
    def select_image(self):
        """选择图像文件"""
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图片", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        
        if file_path:
            self.current_image_path = file_path
            self.display_image(file_path)
            self.detect_btn.setEnabled(True)
    
    def display_image(self, image_path):
        """显示图像"""
        image = QImage(image_path)
        if image.isNull():
            QMessageBox.information(self, "无法加载图像", f"无法加载图像: {image_path}")
            return
        
        # 缩放图像以适应标签大小
        scaled_image = image.scaled(
            self.image_label.size(),
            Qt.KeepAspectRatio,
            Qt.SmoothTransformation
        )
        
        self.image_label.setPixmap(QPixmap.fromImage(scaled_image))
    
    def detect_fruit(self):
        """检测水果成熟度"""
        if not self.model or not self.current_image_path:
            return
        
        # 禁用按钮
        self.select_image_btn.setEnabled(False)
        self.detect_btn.setEnabled(False)
        
        # 创建检测线程
        confidence_threshold = self.confidence_slider.value() / 100.0
        self.detection_thread = DetectionThread(
            self.model, 
            self.current_image_path,
            self.fruit_type_combo.currentText(),
            confidence_threshold,
            self.image_preprocessing
        )
        
        # 连接信号
        self.detection_thread.update_signal.connect(self.update_detection_results)
        self.detection_thread.finished_signal.connect(self.on_detection_finished)
        self.detection_thread.progress_signal.connect(self.progress_bar.setValue)
        
        # 开始检测
        self.progress_bar.setValue(0)
        self.detection_thread.start()
    
    def update_detection_results(self, processed_image, detections):
        """更新检测结果"""
        self.detections = detections
        
        # 显示处理后的图像
        if processed_image is not None:
            # 转换OpenCV图像为Qt图像
            rgb_image = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            
            # 缩放图像以适应标签大小
            scaled_image = q_image.scaled(
                self.image_label.size(),
                Qt.KeepAspectRatio,
                Qt.SmoothTransformation
            )
            
            self.image_label.setPixmap(QPixmap.fromImage(scaled_image))
        
        # 显示检测结果
        self.display_detection_results(detections)
        
        # 更新图表
        self.visualizer.update_plot(detections)
    
    def display_detection_results(self, detections):
        """显示检测结果文本"""
        if not detections:
            self.results_table.setText("未检测到任何水果")
            return
        
        # 使用全局映射替换中文标签
        result_title = RIPENESS_MAPPING.get("检测结果", "Detection Results")
        fruit_type = RIPENESS_MAPPING.get("水果类型", "Fruit Type")
        confidence = RIPENESS_MAPPING.get("置信度", "Confidence")
        ripeness = RIPENESS_MAPPING.get("成熟度", "Ripeness")
        color_analysis = RIPENESS_MAPPING.get("颜色分析", "Color Analysis")
        statistics_title = RIPENESS_MAPPING.get("统计信息", "Statistics")
        
        result_text = f"<h3>{result_title}</h3>"
        result_text += "<table border='1' cellspacing='0' cellpadding='5'>"
        result_text += f"<tr><th>{fruit_type}</th><th>{confidence}</th><th>{ripeness}</th><th>{color_analysis}</th></tr>"
        
        for detection in detections:
            # 使用全局映射替换中文标签
            fruit_name = RIPENESS_MAPPING.get(detection["fruit_name"], detection["fruit_name"])
            conf_value = f"{detection['confidence']:.2f}"
            ripe_value = RIPENESS_MAPPING.get(detection["ripeness"], detection["ripeness"])
            
            # 根据成熟度设置颜色
            color_map = {
                "成熟": "green",
                "半成熟": "orange",
                "未成熟": "blue",
                "过熟": "red",
                "Ripe": "green",
                "Half-ripe": "orange",
                "Unripe": "blue",
                "Overripe": "red"
            }
            ripe_color = color_map.get(ripe_value, "black")
            
            # 颜色信息
            color_info = ", ".join([f"{k}: {v}" for k, v in detection["color_info"].items()])
            
            result_text += f"<tr><td>{fruit_name}</td><td>{conf_value}</td><td style='color:{ripe_color}'>{ripe_value}</td><td>{color_info}</td></tr>"
        
        result_text += "</table>"
        
        # 添加统计信息
        result_text += f"<h3>{statistics_title}</h3>"
        result_text += "<ul>"
        
        # 统计每种水果的数量
        fruit_counts = {}
        for detection in detections:
            # 使用全局映射替换中文标签
            fruit_name = RIPENESS_MAPPING.get(detection["fruit_name"], detection["fruit_name"])
            if fruit_name in fruit_counts:
                fruit_counts[fruit_name] += 1
            else:
                fruit_counts[fruit_name] = 1
        
        for fruit, count in fruit_counts.items():
            result_text += f"<li>{fruit}: {count}个</li>"
        
        result_text += "</ul>"
        
        self.results_table.setText(result_text)
    
    def on_detection_finished(self):
        """检测完成后的处理"""
        # 启用按钮
        self.select_image_btn.setEnabled(True)
        self.detect_btn.setEnabled(True)


if __name__ == "__main__":
    # 设置高DPI支持
    QApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
    QApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)
    
    app = QApplication(sys.argv)
    
    # 设置全局字体
    font = QFont()
    font.setFamily("SimHei")  # 使用黑体
    font.setPointSize(10)
    app.setFont(font)
    
    window = FruitRipenessDetector()
    window.show()
    
    sys.exit(app.exec_())